import os
import pickle
import pandas as pd
import numpy as np


def read_rating_data(path='H:/desktop/ml-1m/ml-1m/ratings.dat', train_rate=1.):
    '''
    读入ratings数据,设置训练集比例返回划分后的训练集和测试集
    :param path: ratings 的文件路径位置
    :param train_rate:  训练集的大小比例
    :return: train_list,test_list集合
    '''
    train_list = list()
    test_list = list()
    data = pd.read_table(path, sep='::', header=None, engine='python')
    data = np.array(data)
    data = list(data)
    p = [train_rate, 1 - train_rate]
    for user, movie, rating, _ in data:
        value = np.random.choice([0, 1], p=p)
        if value == 0:
            train_list.append((user, movie, rating))
        else:
            test_list.append((user, movie, rating))

    return train_list, test_list


def all_items(path='H:/desktop/ml-1m/ml-1m/ratings.dat'):
    '''
    返回全部的movie_id集合
    :param path: ratings文件的路径
    :return: ,movie_id集合
    '''
    item = set()
    data = pd.read_table(path, sep='::', header=None, engine='python')
    data = list(np.array(data))
    for user, movie, rating, _ in data:
        item.add(movie)
    return item


def save_file(filepath, data):
    """
        保存数据
        :param filepath:    保存路径
        :param data:    要保存的数据
    """
    parent_path = filepath[: filepath.rfind("/")]

    if not os.path.exists(parent_path):
        os.mkdir(parent_path)
    with open(filepath, "wb") as f:
        pickle.dump(data, f)


def load_file(filepath):
    """载入二进制数据"""
    with open(filepath, "rb") as f:
        data = pickle.load(f)
    return data


def open_text(filename, skip_row=0):
    """打开文本文件
    :param filename: str
        文件名
    :param skip_row: int
         需要跳过的行数
    :return generator
        生成每一行的文本
    """
    with open(filename, "r", encoding="utf-8") as f:
        for i, line in enumerate(f):
            if i < skip_row:
                continue
            yield line


def save_to_txt(data, path='./output.txt'):
    '''
    将运行后的结果保存,追加到output.txt中
    :param data: 每一次运行结束后的数据
    :param path: output.txt文件路径,若不在则再当前文件夹下自动创建
    :return: 生成一个output.txt包括每次运行的推荐结果
    '''
    if os.path.exists(path):
        d = pd.DataFrame.from_dict(data)

        d = d.transpose()

        d.to_csv('./tmp.txt', sep='\t',index=True)
        f_append = open(path, 'a+')
        with open("./tmp.txt",'r') as f:
            f_append.write('#########################################################')
            f_append.write('\n')
            for i in f:
                f_append.write(i)
        os.remove('./tmp.txt')
    else:
        d = pd.DataFrame.from_dict(data)
        d = d.transpose()

        d.to_csv(path, sep='\t',index=True)